package com.usoft.thrift.benchmark;

import com.usoft.thrift.benchmark.service.MyHelloService;
import org.apache.thrift.TProcessorFactory;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TJSONProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.protocol.TTupleProtocol;
import org.apache.thrift.server.THsHaServer;
import org.apache.thrift.server.TNonblockingServer;
import org.apache.thrift.server.TServer;
import org.apache.thrift.server.TSimpleServer;
import org.apache.thrift.server.TThreadPoolServer;
import org.apache.thrift.server.TThreadedSelectorServer;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TNonblockingServerSocket;
import org.apache.thrift.transport.TNonblockingServerTransport;
import org.apache.thrift.transport.TSSLTransportFactory;
import org.apache.thrift.transport.TServerSocket;
import org.apache.thrift.transport.TServerTransport;
import org.apache.thrift.transport.TTransportException;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;

/**
 * 运行参数:
 * --address=127.0.0.1:9998
 * --server=T_HS_HA_SERVER
 * --transport=T_NONBLOCKING_SERVER_SOCKET
 * --protocol=T_COMPACT_PROTOCOL
 * ===========
 * --address=127.0.0.1:9998 --server=T_THREAD_POOL_SERVER
 * --transport=T_SERVER_SOCKET --protocol=T_COMPACT_PROTOCOL
 * Created by xinxingegeya on 16/2/25.
 */
public class ThriftQpsServer {

    public static void main(String... args) throws Exception {
        new ThriftQpsServer().run(args);
    }

    private void run(String[] args) throws TTransportException {

        ServerConfiguration.Builder configBuilder = ServerConfiguration
            .newBuilder();
        ServerConfiguration config;
        try {
            config = configBuilder.build(args);
        } catch (Exception e) {
            System.out.println(e.getMessage());
            configBuilder.printUsage();
            return;
        }

        final TServer server = newServer(config);
        System.out.println("QPS Server starting on port:"
            + ((InetSocketAddress) config.address).getPort());
        System.out.println("server=" + config.server + ",transport="
            + config.transport + ",protocol=" + config.protocolType);
        server.serve();

        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                try {
                    System.out.println("QPS Server shutting down");
                    server.stop();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });

    }

    private TServer newServer(ServerConfiguration config)
            throws TTransportException {

        TServerTransport serverTransport = newTransport(config);
        TProtocolFactory protocolFactory = newProtocol(config);
        TProcessorFactory processorFactory = newProcessor();

        TServer server;

        switch (config.server) {
            case T_SIMPLE_SERVER:
                TSimpleServer.Args simpleServerArgs = new TSimpleServer.Args(
                    serverTransport);
                simpleServerArgs.processorFactory(processorFactory);
                simpleServerArgs.protocolFactory(protocolFactory);
                server = new TSimpleServer(simpleServerArgs);
                break;
            case T_NONBLOCKING_SERVER:
                TNonblockingServer.Args nonBlockingServerArgs = new TNonblockingServer.Args(
                    (TNonblockingServerTransport) serverTransport);
                nonBlockingServerArgs.processorFactory(processorFactory);
                nonBlockingServerArgs.protocolFactory(protocolFactory);

                /**
                 * To use this server, you MUST use a TFramedTransport at the
                 * outermost
                 * transport, otherwise this server will be unable to determine
                 * when a whole
                 * method call has been read off the wire. Clients must also use
                 * TFramedTransport.
                 */
                nonBlockingServerArgs
                    .transportFactory(new TFramedTransport.Factory());
                server = new TNonblockingServer(nonBlockingServerArgs);
                break;
            case T_HS_HA_SERVER:
                THsHaServer.Args hshaServerArgs = new THsHaServer.Args(
                    (TNonblockingServerTransport) serverTransport);
                hshaServerArgs.processorFactory(processorFactory);
                hshaServerArgs.protocolFactory(protocolFactory);
                /**
                 * Like TNonblockingServer, it relies on the use of
                 * TFramedTransport.
                 */
                hshaServerArgs.transportFactory(new TFramedTransport.Factory());
                hshaServerArgs.maxWorkerThreads(config.maxWorkerThreads)
                    .minWorkerThreads(config.minWorkerThreads)
                    .stopTimeoutVal(config.stopTimeoutVal)
                    .stopTimeoutUnit(config.stopTimeoutUnit);
                server = new THsHaServer(hshaServerArgs);
                break;
            case T_THREAD_SELECTOR_SERVER:
                TThreadedSelectorServer.Args threadSelectorServerArgs = new TThreadedSelectorServer.Args(
                    (TNonblockingServerTransport) serverTransport);
                threadSelectorServerArgs.processorFactory(processorFactory);
                threadSelectorServerArgs.protocolFactory(protocolFactory);
                /**
                 * Like TNonblockingServer, it relies on the use of
                 * TFramedTransport.
                 */
                threadSelectorServerArgs
                    .transportFactory(new TFramedTransport.Factory());
                threadSelectorServerArgs.selectorThreads(config.selectorThreads)
                    .stopTimeoutVal(config.stopTimeoutVal)
                    .stopTimeoutUnit(config.stopTimeoutUnit)
                    .workerThreads(config.workerThreads);

                server = new TThreadedSelectorServer(threadSelectorServerArgs);
                break;

            case T_THREAD_POOL_SERVER:
                TThreadPoolServer.Args threadPoolServerArgs = new TThreadPoolServer.Args(
                    serverTransport);
                threadPoolServerArgs.processorFactory(processorFactory);
                threadPoolServerArgs.protocolFactory(protocolFactory);
                threadPoolServerArgs.maxWorkerThreads(config.maxWorkerThreads)
                    .minWorkerThreads(config.minWorkerThreads)
                    .stopTimeoutVal(config.stopTimeoutVal);
                server = new TThreadPoolServer(threadPoolServerArgs);
                break;

            default:
                server = null;

        }
        return server;
    }

    private TServerTransport newTransport(ServerConfiguration config) {
        TServerTransport serverTransport = null;
        if (config.ssl) {
            TSSLTransportFactory.TSSLTransportParameters params = new TSSLTransportFactory.TSSLTransportParameters();
            params.setKeyStore("path to keystore.jks", "keystore.jks password");
            try {
                serverTransport = TSSLTransportFactory.getServerSocket(7911,
                    10000, InetAddress.getByName("localhost"), params);

            } catch (TTransportException e) {
                e.printStackTrace();
            } catch (UnknownHostException e) {
                e.printStackTrace();
            }
        } else {
            /**
             * TServerTransport的两种实现类:
             * 1)TServerSocket
             * 2)TNonblockingServerSocket
             */
            switch (config.transport) {
                case T_SERVER_SOCKET:
                    try {
                        serverTransport = new TServerSocket(
                            (InetSocketAddress) config.address);
                    } catch (TTransportException e) {
                        e.printStackTrace();
                    }
                    break;
                case T_NONBLOCKING_SERVER_SOCKET:
                    try {
                        serverTransport = new TNonblockingServerSocket(
                            (InetSocketAddress) config.address);
                    } catch (TTransportException e) {
                        e.printStackTrace();
                    }
                    break;
                default:
                    throw new RuntimeException("exception");
            }
        }
        return serverTransport;
    }

    private TProcessorFactory newProcessor() {
        HelloService.Processor<HelloService.Iface> processor = new HelloService.Processor<HelloService.Iface>(
            new MyHelloService());
        return new TProcessorFactory(processor);
    }

    private TProtocolFactory newProtocol(ServerConfiguration config) {
        switch (config.protocolType) {
            case T_BINARY_PROTOCOL:
                return new TBinaryProtocol.Factory();
            case T_COMPACT_PROTOCOL:
                return new TCompactProtocol.Factory();
            case T_TUPLE_PROTOCOL:
                return new TTupleProtocol.Factory();
            case T_JSON_PROTOCOL:
                return new TJSONProtocol.Factory();
            default:
                throw new RuntimeException("not support protocol");
        }
    }

}
